from typing import Tuple
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from time import time
import os

from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import KFold, cross_val_score, train_test_split
from sklearn.naive_bayes import GaussianNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
from sklearn.tree import DecisionTreeClassifier

class AlgoCmp:
    def __init__(self, config):
        self.config = config
        self.config["dataset_name"] = config["dataset_path"].rsplit("/", 1)[-1].rsplit(".")[0]
        self.config["report_dir"] = os.path.join(self.config["output_dir"], self.config["dataset_name"])

        if not os.path.exists(self.config["report_dir"]):
            os.makedirs(self.config["report_dir"])

        self.seed = 7
    
        self.models = {
            "LR": LogisticRegression(max_iter=1000),
            "LDA": LinearDiscriminantAnalysis(),
            "KNN": KNeighborsClassifier(),
            "CART": DecisionTreeClassifier(),
            # "NB": GaussianNB(),
            # "SVM": SVC()
        }

    def compare(self):
        print("Algorithm Comparison")
        results = []
        num_of_folds = 5
        scoring = "accuracy"

        X, y = self._load_data(self.config["dataset_path"])

        for key in self.models:
            t0 = time()
            kfold = KFold(n_splits=num_of_folds, shuffle=True, random_state=self.seed)
            cv_results = cross_val_score(self.models[key], X, y, cv=kfold, scoring=scoring)
            results.append(cv_results)
            t1 = time()
            print("-" * 50)
            print("Duration: {:.2f}".format(t1 - t0))
            print("{}: {:.2f} ({:.2f})".format(self.models[key], cv_results.mean(), cv_results.std()))
            print("-" * 50)
            print()

        self._visualize(results, self.models, "algo_cmp", title="Algorithm Comparison")

    def compare_scaled(self):
        print("Scaled Algorithm Comparison")
        results = []
        num_of_folds = 5
        scoring = "accuracy"

        X, y= self._load_data(self.config["dataset_path"])

        pipelines = {}

        for key in self.models:
            pipelines["scaler_{}".format(key)] = Pipeline([("Scaler", StandardScaler()), (key, self.models[key])])

        for key in pipelines:
            t0 = time()
            kfold = KFold(n_splits=num_of_folds, shuffle=True, random_state=self.seed)
            cv_results = cross_val_score(pipelines[key], X, y, cv=kfold, scoring=scoring)
            results.append(cv_results)
            t1 = time()
            print("-" * 50)
            print("Duration: {:.2f}".format(t1 - t0))
            print("{}: {:.2f} ({:.2f})".format(pipelines[key], cv_results.mean(), cv_results.std()))
            print("-" * 50)
            print()
        
        self._visualize(results, self.models, "scaled_algo_cmp",title="Scaled Algorithm Comparison")


    def _load_data(self,
        dataset_path:str,
        num_of_domains:int=100,
        num_of_files_per_domain:int=100) -> Tuple[np.ndarray, np.ndarray]:
        """Load data.

        Args:
            dataset_path: Path to dataset.
            num_of_domains: Number of domains need to be loaded. 0 means all.
            num_of_files_per_domain: The number of pcap files per domain folder.

        Retuens:
            X, y
        """
        if num_of_domains == 0:
            df = pd.read_csv(dataset_path)
        else:
            df = pd.read_csv(dataset_path, nrows=num_of_domains*num_of_files_per_domain)

        dataset = df.values

        X = dataset[:, :-1]
        y = dataset[:, -1]

        # validation_size = 0.2
        # seed = 7

        # X_train, X_validation, Y_train, Y_validation = train_test_split(
        #     X, Y, test_size=validation_size, random_state=seed)
        
        # return X_train, X_validation, Y_train, Y_validation

        return X, y

    def _visualize(self, results:list, models:dict, fig_name:str, title:str="") -> str:
        """Generate Box-plot.
        """
        fig_path = os.path.join(self.config["report_dir"], fig_name+".svg")

        fig = plt.figure()
        ax = fig.add_subplot(111)
        fig.suptitle(title)
        plt.boxplot(results)
        ax.set_xticklabels(models.keys())
        # plt.show()
        plt.savefig(fig_path, format="svg")

        return fig_path

def main():
    config = {
        "dataset_path": "data/dataset/dataset-nfstream-20220304210805-shuffled.csv",
        "output_dir": "data/dataset_report"
    }
    
    ac = AlgoCmp(config)
    # ac.compare()
    ac.compare_scaled()
    
if __name__ == "__main__":
    main()